import pandas as pd
import numpy as np
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.ensemble import RandomForestRegressor
import matplotlib.pyplot as plt
import seaborn as sns
import json
import sys
import os
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def main(selected_data_csv_path):
    try:
        print(f"开始数据插补，数据文件: {selected_data_csv_path}")

        # 读取21个字段数据
        df = pd.read_csv(selected_data_csv_path)
        print(f"原始数据形状: {df.shape}")
        print(f"原始数据列: {list(df.columns)}")

        # 记录原始数据信息
        original_shape = df.shape
        original_missing = df.isnull().sum()
        original_missing_total = original_missing.sum()

        print(f"原始缺失值统计:")
        for col, missing_count in original_missing.items():
            if missing_count > 0:
                print(f"  {col}: {missing_count} ({missing_count/len(df)*100:.2f}%)")

        # 创建带缺失值的数据副本（如果原数据没有缺失值，则模拟一些）
        df_with_missing = df.copy()

        if original_missing_total == 0:
            print("原始数据没有缺失值，模拟添加10%的缺失值进行演示...")
            np.random.seed(42)
            for col in df.columns:
                if df[col].dtype in ['float64', 'int64']:
                    # 随机设置10%的值为缺失值
                    mask = np.random.rand(len(df)) < 0.1
                    df_with_missing.loc[mask, col] = np.nan

        missing_before = df_with_missing.isnull().sum()
        missing_before_total = missing_before.sum()

        print(f"插补前缺失值统计:")
        for col, missing_count in missing_before.items():
            if missing_count > 0:
                print(f"  {col}: {missing_count} ({missing_count/len(df_with_missing)*100:.2f}%)")

        # 配置IterativeImputer
        imputer = IterativeImputer(
            estimator=RandomForestRegressor(n_estimators=10, random_state=42),
            random_state=42,
            max_iter=10,
            verbose=1
        )

        print("开始执行IterativeImputer插补...")

        # 执行插补
        imputed_array = imputer.fit_transform(df_with_missing)
        imputed_df = pd.DataFrame(imputed_array, columns=df.columns)

        print("插补完成!")

        # 检查插补后的缺失值
        missing_after = imputed_df.isnull().sum()
        missing_after_total = missing_after.sum()

        print(f"插补后缺失值统计:")
        for col, missing_count in missing_after.items():
            if missing_count > 0:
                print(f"  {col}: {missing_count}")

        if missing_after_total == 0:
            print("所有缺失值已成功插补!")

        # 保存插补后的数据
        os.makedirs('temp/output', exist_ok=True)
        output_path = 'temp/output/imputed_data.csv'
        imputed_df.to_csv(output_path, index=False)

        print(f"插补后数据已保存到: {output_path}")

        # 创建插补前后对比图
        comparison_plot_path = create_imputation_comparison_plot(
            df_with_missing, imputed_df, missing_before)

        # 计算插补质量指标
        quality_metrics = calculate_imputation_quality(df_with_missing, imputed_df)

        # 构建结果
        result = {
            'imputed_data_path': output_path,
            'original_shape': list(original_shape),
            'imputed_shape': list(imputed_df.shape),
            'missing_count_before': missing_before.to_dict(),
            'missing_count_after': missing_after.to_dict(),
            'missing_total_before': int(missing_before_total),
            'missing_total_after': int(missing_after_total),
            'imputation_successful': missing_after_total == 0,
            'comparison_plot_path': comparison_plot_path,
            'quality_metrics': quality_metrics,
            'columns': list(imputed_df.columns),
            'data_preview': imputed_df.head(5).to_dict('records'),
            'status': 'success'
        }



        py_result = to_python_types(result)

        print(json.dumps(py_result))

    except Exception as e:
        error_result = {
            'status': 'error',
            'message': f'数据插补失败: {str(e)}'
        }
        print(json.dumps(error_result))



# 手动转换整个 dict
def to_python_types(obj):
    if isinstance(obj, np.generic):
        return obj.item()           # .item() 会把 numpy 类型变成原生 Python
    elif isinstance(obj, dict):
        return {k: to_python_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [to_python_types(v) for v in obj]
    else:
        return obj


def create_imputation_comparison_plot(df_before, df_after, missing_counts):
    """创建插补前后对比图，只保存 QYB 字段的对比图"""
    try:
        field = "qyb"
        # 确保输出目录存在
        os.makedirs('temp/output', exist_ok=True)

        plt.figure(figsize=(10, 6))

        # 原始（插补前）分布
        sns.histplot(df_before[field].dropna(), kde=True,
                     color='red', alpha=0.5, label='插补前', bins=30)
        # 插补后分布
        sns.histplot(df_after[field].dropna(), kde=True,
                     color='blue', alpha=0.5, label='插补后', bins=30)

        plt.title(f'{field} 插补前后数据分布对比')
        plt.xlabel(field)
        plt.ylabel('频率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()

        comparison_plot_path = 'temp/output/QYB_imputation_comparison.png'
        plt.savefig(comparison_plot_path, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"插补对比图已保存到: {comparison_plot_path}")
        return comparison_plot_path

    except Exception as e:
        print(f"创建对比图失败: {e}")
        return None




def calculate_imputation_quality(df_before, df_after):
    """计算插补质量指标"""
    quality_metrics = {}

    try:
        # 计算每个字段的基本统计信息变化
        for col in df_before.columns:
            if df_before[col].dtype in ['float64', 'int64']:
                before_stats = df_before[col].describe()
                after_stats = df_after[col].describe()

                quality_metrics[col] = {
                    'mean_change': float(after_stats['mean'] - before_stats['mean']),
                    'std_change': float(after_stats['std'] - before_stats['std']),
                    'missing_count': int(df_before[col].isnull().sum())
                }

    except Exception as e:
        print(f"计算质量指标失败: {e}")
        quality_metrics = {}

    return quality_metrics

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(json.dumps({
            'status': 'error',
            'message': '参数错误，需要提供CSV文件路径'
        }))
    else:
        main(sys.argv[1])
